feat(data): Add AseDataLoader utility with lazy loading for ase - #15
Conversation
|
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
muhrin
left a comment
muhrin
left a comment
There was a problem hiding this comment.
Looks good, just a few changes to be made. Also, if you could turn on pre-commit checks that would be good, as some of them are failing for this PR. You can do this using:
pre-commit install
in the terminal. In case the pre-commit package is, itself, not installed, just go to your tensorial folder and do:
pip install -e .[dev]
| self._read_kwargs: Final[dict[str, Any]] = self._init_kwargs(limit, read_kwargs) | ||
|
|
||
| try: | ||
| import ase.io |
There was a problem hiding this comment.
Instead of doing this can you use the lazy import helper that I left in the chat. Basically, as a general rule it's not a good idea to do imports anywhere other than the top of the module file because it harms the readability of functions, and if we have everything at the top of the module it's easy to see all the direct dependencies of this module.
Here's the example:
import importlib.util
import sys
def lazy_import(name: str):
"""Lazily import a module using the standard library."""
spec = importlib.util.find_spec(name)
if spec is None:
# Optimization: if it's completely missing, we can fail early
# or return a dummy object.
pass
loader = importlib.util.LazyLoader(spec.loader)
module = importlib.util.module_from_spec(spec)
spec.loader = loader
sys.modules[name] = module
return module
# Usage at top of file
np = lazy_import("numpy")
| @@ -0,0 +1,77 @@ | |||
| """Module for loading ase.Atoms objects as graphs""" | |||
|
|
|||
| import collections.abc | |||
There was a problem hiding this comment.
This you can change to from collections.abc import Sequence
| __all__ = ("AseDataLoader",) | ||
|
|
||
|
|
||
| class AseDataLoader(collections.abc.Sequence[jraph.GraphsTuple]): |
There was a problem hiding this comment.
...and here the parent class can just be Sequence[jraph.GraphsTuple]
| from typing import TYPE_CHECKING, Any, Final | ||
|
|
||
| import jraph | ||
| from tensorial import gcnn |
There was a problem hiding this comment.
You cannot import gcnn this way, because you are 'inside' gcnn and so this creates a cyclic import dependency.
Instead what you can do is from .. import atomic which is now relative (and specific), hence avoiding the cyclic import.
| entry = self._data[item] | ||
| if self._to_graphs and not isinstance(entry, jraph.GraphsTuple): | ||
| # Lazily convert the first time | ||
| entry = gcnn.atomic.graph_from_ase(entry, **self._to_graphs) |
There was a problem hiding this comment.
...and here you just remove gcnn and call atomic.graph_from... directly.
…er utility with lazy loading
bc1e48e to
8f60ddf
Compare
|
Thanks! |
3 participants
Lazy loader for ASE Atoms with optional conversion to jraph.GraphsTuple.